import os
import sys
import subprocess
import importlib
import platform

def run_command(command):
    """运行命令并处理可能的错误"""
    try:
        subprocess.check_call(command)
        return True
    except subprocess.CalledProcessError:
        print(f"命令执行失败: {' '.join(command)}")
        return False

def check_and_install_dependencies():
    """检查并安装所需的依赖"""
    print("正在检查并安装必要的依赖...")
    
    # 检查 PyTorch 和 Torchvision
    try:
        import torch
        import torchvision
        print(f"PyTorch 版本: {torch.__version__}")
        print(f"Torchvision 版本: {torchvision.__version__}")
    except ImportError:
        print("正在安装 PyTorch 和 Torchvision...")
        # 根据系统选择合适的安装命令
        if platform.system() == "Windows":
            success = run_command([
                sys.executable, "-m", "pip", "install", 
                "torch", "torchvision", "--index-url", "https://download.pytorch.org/whl/cu118"
            ])
        else:
            success = run_command([
                sys.executable, "-m", "pip", "install", "torch", "torchvision"
            ])
        
        if not success:
            print("PyTorch 安装失败，请手动安装。")
            print("可以访问 https://pytorch.org/get-started/locally/ 获取安装指南。")
            return False
    
    # 安装 opencv-python
    try:
        import cv2
        print(f"OpenCV 版本: {cv2.__version__}")
    except ImportError:
        print("正在安装 OpenCV...")
        if not run_command([sys.executable, "-m", "pip", "install", "opencv-python"]):
            print("OpenCV 安装失败，请手动安装。")
            return False
    
    # 安装 paddleocr 和 paddlepaddle
    try:
        import paddleocr
        print("PaddleOCR 已安装")
    except ImportError:
        print("正在安装 PaddleOCR...")
        if not run_command([sys.executable, "-m", "pip", "install", "paddleocr"]):
            print("PaddleOCR 安装失败，请手动安装。")
            return False
    
    try:
        import paddle
        print(f"PaddlePaddle 版本: {paddle.__version__}")
    except ImportError:
        print("正在安装 PaddlePaddle...")
        if not run_command([sys.executable, "-m", "pip", "install", "paddlepaddle"]):
            print("PaddlePaddle 安装失败，请手动安装。")
            return False
    
    # 尝试使用简化版的布局分析
    print("\n由于 Detectron2 安装困难，我们将使用简化版的布局分析方法。")
    return True

def update_imageocr_script():
    """更新 imageocr.py 脚本以使用简化的布局分析"""
    try:
        with open('imageocr.py', 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 移除 layoutparser 相关代码
        updated_content = content.replace(
            "import layoutparser as lp", 
            "# import layoutparser as lp  # 暂时移除，使用简化版布局分析"
        )
        
        # 替换布局分析模型初始化部分
        layout_model_init = """
    def __init__(self, lang='ch'):
        # 初始化OCR引擎
        self.ocr = PaddleOCR(use_angle_cls=True, lang=lang)
        
        # 初始化布局分析模型
        self.layout_model = lp.Detectron2LayoutModel(
            'lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config',
            extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.5],
            label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}
        )"""
        
        simplified_model_init = """
    def __init__(self, lang='ch'):
        # 初始化OCR引擎
        self.ocr = PaddleOCR(use_angle_cls=True, lang=lang)
        
        # 使用简化版布局分析
        print("使用简化版布局分析方法")"""
        
        updated_content = updated_content.replace(layout_model_init, simplified_model_init)
        
        # 替换布局分析方法
        process_page_method = """
    def process_single_image(self, img_path):
        \"\"\"处理单个图像，提取文本和布局\"\"\"
        # 读取图像
        image = cv2.imread(img_path)
        if image is None:
            raise ValueError(f"无法读取图像: {img_path}")
        
        # 预处理图像（可选）
        # image = self.preprocess_image(image)
        
        # OCR识别
        ocr_result = self.ocr.ocr(img_path, cls=True)
        
        # 布局分析
        layout_image = lp.load_image(img_path)
        layout = self.layout_model.detect(layout_image)"""
        
        simplified_process_page = """
    def process_single_image(self, img_path):
        \"\"\"处理单个图像，提取文本和布局\"\"\"
        # 读取图像
        image = cv2.imread(img_path)
        if image is None:
            raise ValueError(f"无法读取图像: {img_path}")
        
        # 预处理图像（可选）
        # image = self.preprocess_image(image)
        
        # OCR识别
        ocr_result = self.ocr.ocr(img_path, cls=True)
        
        # 简化版布局分析 - 使用基于OCR结果的启发式方法
        height, width = image.shape[:2]"""
        
        updated_content = updated_content.replace(process_page_method, simplified_process_page)
        
        # 替换布局处理部分
        layout_processing = """
        # 处理布局元素
        for block in layout:
            block_coords = block.coordinates
            x1, y1, x2, y2 = [int(coord) for coord in [block_coords[0], block_coords[1], block_coords[2], block_coords[3]]]
            
            # 提取该区域内的OCR结果
            block_texts = []
            for line in ocr_result[0]:
                bbox = line[0]
                text = line[1][0]
                confidence = line[1][1]
                
                # 检查文本框是否在当前布局块内
                text_x1, text_y1 = bbox[0][0], bbox[0][1]
                text_x2, text_y2 = bbox[2][0], bbox[2][1]
                
                if (x1 <= text_x1 <= x2 and y1 <= text_y1 <= y2) or \\
                   (x1 <= text_x2 <= x2 and y1 <= text_y2 <= y2):
                    block_texts.append({
                        "text": text,
                        "confidence": float(confidence),
                        "bbox": [[float(p[0]), float(p[1])] for p in bbox]
                    })
            
            # 根据布局类型处理
            if block.type == "Table":
                # 对于表格，尝试提取行和列
                table_content = self.extract_table_structure(image[y1:y2, x1:x2], block_texts)
                result["blocks"].append({
                    "type": "table",
                    "coordinates": [x1, y1, x2, y2],
                    "content": table_content
                })
            else:
                # 对于其他类型的块
                result["blocks"].append({
                    "type": block.type.lower(),
                    "coordinates": [x1, y1, x2, y2],
                    "content": block_texts
                })"""
        
        simplified_layout_processing = """
        # 简化版布局分析 - 基于OCR结果进行分组
        result["blocks"] = []
        
        if ocr_result[0]:
            # 所有文本框
            all_texts = []
            for line in ocr_result[0]:
                bbox = line[0]
                text = line[1][0]
                confidence = line[1][1]
                all_texts.append({
                    "text": text,
                    "confidence": float(confidence),
                    "bbox": [[float(p[0]), float(p[1])] for p in bbox]
                })
            
            # 尝试检测表格区域 (简化版)
            # 这里使用启发式方法：检查文本框的排列是否规则
            table_candidates = self.detect_tables_heuristic(all_texts, width, height)
            
            # 添加检测到的表格
            for table in table_candidates:
                table_texts = []
                for text in all_texts:
                    # 检查文本是否在表格区域内
                    text_x1 = text["bbox"][0][0]
                    text_y1 = text["bbox"][0][1]
                    if (table[0] <= text_x1 <= table[2] and table[1] <= text_y1 <= table[3]):
                        table_texts.append(text)
                
                if table_texts:
                    table_content = self.extract_table_structure(
                        image[int(table[1]):int(table[3]), int(table[0]):int(table[2])], 
                        table_texts
                    )
                    result["blocks"].append({
                        "type": "table",
                        "coordinates": [int(coord) for coord in table],
                        "content": table_content
                    })
            
            # 剩余文本作为普通文本块
            remaining_texts = []
            for text in all_texts:
                is_in_table = False
                for table in table_candidates:
                    text_x1 = text["bbox"][0][0]
                    text_y1 = text["bbox"][0][1]
                    if (table[0] <= text_x1 <= table[2] and table[1] <= text_y1 <= table[3]):
                        is_in_table = True
                        break
                
                if not is_in_table:
                    remaining_texts.append(text)
            
            if remaining_texts:
                # 按Y坐标分组为段落
                paragraphs = self.group_texts_into_paragraphs(remaining_texts)
                
                for i, para_texts in enumerate(paragraphs):
                    # 计算段落边界
                    min_x = min(text["bbox"][0][0] for text in para_texts)
                    min_y = min(text["bbox"][0][1] for text in para_texts)
                    max_x = max(text["bbox"][2][0] for text in para_texts)
                    max_y = max(text["bbox"][2][1] for text in para_texts)
                    
                    # 判断是否为标题 (简单启发式：短且位于顶部)
                    is_title = (len(para_texts) <= 2 and min_y < height * 0.2)
                    
                    result["blocks"].append({
                        "type": "title" if is_title else "text",
                        "coordinates": [int(min_x), int(min_y), int(max_x), int(max_y)],
                        "content": para_texts
                    })"""
        
        updated_content = updated_content.replace(layout_processing, simplified_layout_processing)
        
        # 添加新的辅助方法
        new_methods = """
    def detect_tables_heuristic(self, texts, width, height):
        \"\"\"使用启发式方法检测表格
        
        基本思路：
        1. 寻找排列整齐的文本行
        2. 检查这些行是否有相似的宽度和间距
        \"\"\"
        if not texts:
            return []
            
        # 按Y坐标排序
        texts_by_y = sorted(texts, key=lambda x: x["bbox"][0][1])
        
        # 分组为行
        rows = []
        current_row = [texts_by_y[0]]
        row_y = texts_by_y[0]["bbox"][0][1]
        
        for text in texts_by_y[1:]:
            text_y = text["bbox"][0][1]
            # 如果Y坐标相差不大，认为是同一行
            if abs(text_y - row_y) < 20:  # 阈值可调整
                current_row.append(text)
            else:
                # 按X坐标排序当前行
                current_row.sort(key=lambda x: x["bbox"][0][0])
                rows.append(current_row)
                # 开始新行
                current_row = [text]
                row_y = text_y
        
        # 处理最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0][0])
            rows.append(current_row)
        
        # 检测表格特征
        table_candidates = []
        
        # 至少需要3行才考虑是表格
        if len(rows) >= 3:
            # 检查行是否有相似的列数和间距
            row_lengths = [len(row) for row in rows]
            
            # 如果多行有相同的列数，可能是表格
            for i in range(len(rows) - 2):
                if row_lengths[i] >= 2 and row_lengths[i] == row_lengths[i+1] == row_lengths[i+2]:
                    # 找到可能的表格，计算边界
                    table_rows = rows[i:i+3]
                    min_x = min(text["bbox"][0][0] for row in table_rows for text in row)
                    min_y = min(text["bbox"][0][1] for row in table_rows for text in row)
                    max_x = max(text["bbox"][2][0] for row in table_rows for text in row)
                    max_y = max(text["bbox"][2][1] for row in table_rows for text in row)
                    
                    # 扩展表格边界，包含后续相似行
                    j = i + 3
                    while j < len(rows) and row_lengths[j] == row_lengths[i]:
                        for text in rows[j]:
                            min_x = min(min_x, text["bbox"][0][0])
                            min_y = min(min_y, text["bbox"][0][1])
                            max_x = max(max_x, text["bbox"][2][0])
                            max_y = max(max_y, text["bbox"][2][1])
                        j += 1
                    
                    # 添加边距
                    padding = 10
                    table_candidates.append([
                        max(0, min_x - padding),
                        max(0, min_y - padding),
                        min(width, max_x + padding),
                        min(height, max_y + padding)
                    ])
                    
                    # 跳过已处理的行
                    i = j - 1
        
        return table_candidates
    
    def group_texts_into_paragraphs(self, texts):
        \"\"\"将文本分组为段落\"\"\"
        if not texts:
            return []
            
        # 按Y坐标排序
        texts_by_y = sorted(texts, key=lambda x: x["bbox"][0][1])
        
        # 分组为段落
        paragraphs = []
        current_para = [texts_by_y[0]]
        para_y = texts_by_y[0]["bbox"][2][1]  # 使用底部Y坐标
        
        for text in texts_by_y[1:]:
            text_y = text["bbox"][0][1]  # 使用顶部Y坐标
            # 如果新文本的顶部与当前段落底部距离较近，认为是同一段落
            if abs(text_y - para_y) < 30:  # 段落间距阈值
                current_para.append(text)
                # 更新段落底部坐标
                para_y = max(para_y, text["bbox"][2][1])
            else:
                # 按X坐标排序当前段落中的文本
                current_para.sort(key=lambda x: x["bbox"][0][0])
                paragraphs.append(current_para)
                # 开始新段落
                current_para = [text]
                para_y = text["bbox"][2][1]
        
        # 处理最后一个段落
        if current_para:
            current_para.sort(key=lambda x: x["bbox"][0][0])
            paragraphs.append(current_para)
            
        return paragraphs"""
        
        # 在类定义的末尾添加新方法
        class_end_index = updated_content.find("if __name__ == \"__main__\":")
        updated_content = updated_content[:class_end_index] + new_methods + "\n\n" + updated_content[class_end_index:]
        
        with open('imageocr_simplified.py', 'w', encoding='utf-8') as f:
            f.write(updated_content)
        
        print("已创建简化版的 imageocr_simplified.py 脚本")
    except Exception as e:
        print(f"更新脚本时出错: {e}")

if __name__ == "__main__":
    print("开始修复 OCR 依赖问题...")
    if check_and_install_dependencies():
        update_imageocr_script()
        print("\n修复完成！请使用以下命令运行简化版OCR脚本:")
        print("python imageocr_simplified.py path/to/your/image.png --output output_folder --visualize")
    else:
        print("\n修复未完成。请手动安装所需依赖。") 